/*
 * Tencent is pleased to support the open source community by making Angel available.
 *
 * Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved.
 *
 * Licensed under the BSD 3-Clause License (the "License"); you may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 * https://opensource.org/licenses/BSD-3-Clause
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is
 * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied. See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.tencent.angel.master.oplog;

import com.tencent.angel.conf.AngelConfiguration;
import com.tencent.angel.exception.InvalidParameterException;
import com.tencent.angel.master.MatrixMetaManager;
import com.tencent.angel.master.app.AMContext;
import com.tencent.angel.master.data.DataSplitter;
import com.tencent.angel.master.ps.ParameterServerManager;
import com.tencent.angel.master.ps.ps.AMParameterServer;
import com.tencent.angel.master.task.AMTaskManager;
import com.tencent.angel.ps.PSAttemptId;
import com.tencent.angel.ps.ParameterServerId;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.PathFilter;
import org.apache.hadoop.service.AbstractService;

/**
 * Application state storage. It write the application state to files every fixed period of time
 */
// TODO: 17/6/26 by zmyer
public class AppStateStorage extends AbstractService {
    private static final Log LOG = LogFactory.getLog(AppStateStorage.class);

    private static final String matrixMetaFileNamePrefix = "matrixmeta";
    private static final String fileNameSeparator = "_";
    private static final String splitFileName = "splits";
    private static final String psMetaFileNamePrefix = "psmeta";
    private static final String taskMetaFileNamePrefix = "taskmeta";
    private static final String tmpPrefix = "_tmp.";

    private final AMContext context;

    /** ps meta and task meta write thread */
    private Thread writter;

    /** base state file directory */
    private final String writeDir;
    private final Path writeDirPath;

    /** ps meta and task meta write interval in milliseconds */
    private final int writeIntervalMS;

    /** whether you want to stop the application state write service */
    private final AtomicBoolean stopped;
    private final FileSystem fs;

    /** data splits file path */
    private final Path splitFilePath;

    /** last matrix meta file path */
    private Path lastMatrixMetaFilePath;

    private final Lock matrixMetaLock;
    private final Lock taskMetaLock;
    private final Lock psMetaLock;

    /**
     * Create a AppStateStorage
     *
     * @param context master context
     * @param writeDir storage file directory
     * @param fs File System
     */
    public AppStateStorage(AMContext context, String writeDir, FileSystem fs) {
        super("app-state-writter");
        LOG.info("writeDir=" + writeDir);
        this.context = context;
        this.writeDir = writeDir;
        this.writeDirPath = new Path(writeDir);
        this.fs = fs;

        splitFilePath = new Path(writeDirPath, splitFileName);
        matrixMetaLock = new ReentrantLock();
        taskMetaLock = new ReentrantLock();
        psMetaLock = new ReentrantLock();

        writeIntervalMS = context.getConf().getInt(
            AngelConfiguration.ANGEL_AM_WRITE_STATE_INTERVAL_MS,
            AngelConfiguration.DEFAULT_ANGEL_AM_WRITE_STATE_INTERVAL_MS);
        this.stopped = new AtomicBoolean(false);
    }

    @Override
    protected void serviceStart() throws Exception {
        startWriter();
        super.serviceStart();
    }

    private void startWriter() {
        writter = new Writter();
        writter.setName("app-state-writter");
        writter.start();
    }

    @Override
    protected void serviceInit(Configuration conf) throws Exception {
        super.serviceInit(conf);
    }

    @Override
    protected void serviceStop() throws Exception {
        if (stopped.getAndSet(true)) {
            return;
        }
        if (writter != null) {
            writter.interrupt();
            try {
                writter.join();
            } catch (InterruptedException ie) {
                LOG.warn("app-state-writter InterruptedException while stopping");
            }
        }

        LOG.info("app-state-writter service stop!");
    }

    /**
     * write training data split meta data to file
     *
     * @param dataSplitter training data split meta storage
     */
    public void writeDataSplits(DataSplitter dataSplitter) throws IOException {
        //write meta data to a temporary file
        String tmpFile = getTmpDataSplitsFile();
        Path tmpPath = new Path(writeDir, tmpFile);
        FSDataOutputStream outputStream = fs.create(tmpPath);
        dataSplitter.serialize(outputStream);
        outputStream.close();

        //rename the temporary file to final file
        if (fs.exists(splitFilePath)) {
            fs.delete(splitFilePath, false);
            fs.rename(tmpPath, splitFilePath);
        } else {
            fs.rename(tmpPath, splitFilePath);
        }
    }

    private String getTmpDataSplitsFile() {
        return tmpPrefix + splitFileName;
    }

    /**
     * load training data meta information from file
     *
     * @return DataSplitter training data split meta storage, if the file does not exist, just
     * return null
     */
    public DataSplitter loadDataSplits() throws IOException, ClassNotFoundException {
        //if the file exists, read from file and deserialize it
        if (fs.exists(splitFilePath)) {
            FSDataInputStream inputStream = fs.open(new Path(writeDir, splitFileName));
            DataSplitter dataSplitter = new DataSplitter(context);
            dataSplitter.deserialize(inputStream);
            return dataSplitter;
        } else {
            LOG.warn("data split file does not exist");
            return null;
        }
    }

    /**
     * write stream into special file object
     */
    private void flushStreamToFile(Path tmpPath, Path matrixFilePath) throws IOException {
        if (fs.exists(matrixFilePath)) {
            fs.delete(matrixFilePath, false);
            fs.rename(tmpPath, matrixFilePath);
        } else {
            fs.rename(tmpPath, matrixFilePath);
            if (lastMatrixMetaFilePath != null) {
                fs.delete(lastMatrixMetaFilePath, false);
            }
        }
    }

    /**
     * write matrix meta to file
     *
     * @param matrixMetaManager matrix meta storage
     * @throws IOException IO Excepiton
     */
    public void writeMatrixMeta(MatrixMetaManager matrixMetaManager) throws IOException {
        try {
            matrixMetaLock.lock();
            String matrixMetaFile = getMatrixMetaFile();
            String tmpFile = getTmpMatrixMetaFile(matrixMetaFile);
            Path tmpPath = new Path(writeDir, tmpFile);
            FSDataOutputStream outputStream = fs.create(tmpPath);
            matrixMetaManager.serialize(outputStream);
            outputStream.close();

            Path matrixFilePath = new Path(writeDir, matrixMetaFile);
            flushStreamToFile(tmpPath, matrixFilePath);
            lastMatrixMetaFilePath = matrixFilePath;

        } finally {
            matrixMetaLock.unlock();
        }
    }

    private String getMatrixMetaFile() {
        return matrixMetaFileNamePrefix + fileNameSeparator + System.nanoTime();
    }

    private String getTmpMatrixMetaFile(String matrixMetaFileName) {
        return tmpPrefix + matrixMetaFileName;
    }

    /**
     * load matrix meta from file
     *
     * @return MatrixMetaManager matrix meta storage
     * @throws IOException, ClassNotFoundException, InvalidParameterException
     */
    public MatrixMetaManager loadMatrixMeta() throws IOException, ClassNotFoundException, InvalidParameterException {
        try {
            matrixMetaLock.lock();
            Path matrixFilePath;
            try {
                matrixFilePath = findFilePathUsePrefix(matrixMetaFileNamePrefix);
            } catch (Exception x) {
                LOG.error("find matrix meta file failed.", x);
                return null;
            }

            if (matrixFilePath == null) {
                return null;
            }

            FSDataInputStream inputStream = fs.open(new Path(writeDir, matrixFilePath));
            MatrixMetaManager matrixMetaManager = new MatrixMetaManager();
            matrixMetaManager.deserialize(inputStream);
            return matrixMetaManager;
        } finally {
            matrixMetaLock.unlock();
        }
    }

    private Path findFilePathUsePrefix(
        final String prefix) throws IOException {
        FileStatus[] files = fs.listStatus(writeDirPath, new PathFilter() {
            @Override
            public boolean accept(Path path) {
                return path.getName().startsWith(prefix);
            }
        });

        return chooseMaxTsFilePath(files);
    }

    private Path chooseMaxTsFilePath(FileStatus[] files) {
        if (files == null || files.length == 0) {
            return null;
        }

        int size = files.length;
        long maxTs = 0;
        int choosedFileIndex = 0;
        for (int i = 0; i < size; i++) {
            String[] nameSplits = files[i].getPath().getName().split(fileNameSeparator);
            if (nameSplits.length == 2) {
                long fileTs = Long.parseLong(nameSplits[1]);
                if (maxTs < fileTs) {
                    choosedFileIndex = i;
                    maxTs = fileTs;
                }
            }
        }

        return files[choosedFileIndex].getPath();
    }

    private class Writter extends Thread {
        @Override
        public void run() {
            while (!stopped.get() && !Thread.interrupted()) {
                long startTs = System.currentTimeMillis();
                //write task meta
                try {
                    writeTaskMeta(context.getTaskManager());
                } catch (IOException e) {
                    if (Thread.interrupted()) {
                        LOG.warn("write task meta is interupted");
                    } else {
                        LOG.error("write task meta file failed.", e);
                    }
                }

                //write ps meta
                try {
                    writePSMeta(context.getParameterServerManager());
                } catch (IOException e) {
                    if (Thread.interrupted()) {
                        LOG.warn("write ps meta is interupted");
                    } else {
                        LOG.error("write ps meta file failed.", e);
                    }
                }

                long useTime = System.currentTimeMillis() - startTs;
                if (useTime < writeIntervalMS) {
                    try {
                        Thread.sleep(writeIntervalMS - useTime);
                    } catch (InterruptedException ignored) {

                    }
                }
            }
        }
    }

    private String getTaskMetaTmpeFile(String taskMetaFileName) {
        return tmpPrefix + taskMetaFileName;
    }

    private String getPSMetaTmpeFile(String psMetaFileName) {
        return tmpPrefix + psMetaFileName;
    }

    private String getTaskMetaFile() {
        return taskMetaFileNamePrefix + fileNameSeparator + System.nanoTime();
    }

    private String getPsMetaFile() {
        return psMetaFileNamePrefix + fileNameSeparator + System.nanoTime();
    }

    /**
     * write task meta to file
     *
     * @param taskManager task meta storage
     * @throws IOException IO Exception that throw
     */
    private void writeTaskMeta(AMTaskManager taskManager) throws IOException {
        try {
            taskMetaLock.lock();
            //generate a temporary file
            String taskMetaFile = getTaskMetaFile();
            String tmpFile = getTaskMetaTmpeFile(taskMetaFile);
            Path tmpPath = new Path(writeDir, tmpFile);

            //write task meta to the temporary file first
            FSDataOutputStream outputStream = fs.create(tmpPath);
            taskManager.serialize(outputStream);
            outputStream.close();

            //rename the temporary file to the final file
            Path taskMetaFilePath = new Path(writeDir, taskMetaFile);
            flushStreamToFile(tmpPath, taskMetaFilePath);

            /* last task meta file path */
        } finally {
            taskMetaLock.unlock();
        }
    }

    /**
     * load task meta from file
     *
     * @return AMTaskManager task meta storage
     * @throws IOException IO Exception that throw
     */
    public AMTaskManager loadTaskMeta() throws IOException {
        try {
            taskMetaLock.lock();
            //find task meta file which has max timestamp
            Path taskMetaFilePath;
            try {
                taskMetaFilePath = findFilePathUsePrefix(taskMetaFileNamePrefix);
            } catch (Exception x) {
                LOG.error("find task meta file failed.", x);
                return null;
            }

            //if the file does not exist, just return null
            if (taskMetaFilePath == null) {
                return null;
            }

            //read task meta from file and deserialize it
            FSDataInputStream inputStream = fs.open(taskMetaFilePath);
            AMTaskManager taskManager = new AMTaskManager();
            taskManager.deserialize(inputStream);
            inputStream.close();
            return taskManager;
        } finally {
            taskMetaLock.unlock();
        }
    }

    /**
     * write ps meta to file
     *
     * @param psManager ps meta storage
     * @throws IOException IO Exception that throw
     */
    private void writePSMeta(ParameterServerManager psManager) throws IOException {
        try {
            psMetaLock.lock();
            //generate a temporary file
            String psMetaFile = getPsMetaFile();
            String tmpFile = getPSMetaTmpeFile(psMetaFile);
            Path tmpPath = new Path(writeDir, tmpFile);
            FSDataOutputStream outputStream = fs.create(tmpPath);

            //write ps meta to the temporary file first.
            Map<ParameterServerId, AMParameterServer> psMap = psManager.getParameterServerMap();
            outputStream.writeInt(psMap.size());
            PSAttemptId attemptId;
            int nextAttemptIndex;
            for (Entry<ParameterServerId, AMParameterServer> entry : psMap.entrySet()) {
                outputStream.writeInt(entry.getKey().getIndex());
                attemptId = entry.getValue().getRunningAttemptId();
                nextAttemptIndex = entry.getValue().getNextAttemptNumber();
                if (attemptId != null) {
                    nextAttemptIndex = attemptId.getIndex();
                }
                outputStream.writeInt(nextAttemptIndex);
            }

            outputStream.close();

            //rename the temporary file to the final file
            Path psMetaFilePath = new Path(writeDir, psMetaFile);
            flushStreamToFile(tmpPath, psMetaFilePath);

        } finally {
            psMetaLock.unlock();
        }
    }

    /**
     * load ps meta from file
     *
     * @return Map<ParameterServerId Integer> psId to attempt index map
     * @throws IOException IO Exception that throw
     */
    public Map<ParameterServerId, Integer> loadPSMeta() throws IOException {
        try {
            psMetaLock.lock();
            //find ps meta file
            Path psMetaFilePath;
            try {
                psMetaFilePath = findFilePathUsePrefix(psMetaFileNamePrefix);
            } catch (Exception x) {
                LOG.error("find ps meta file failed.", x);
                return null;
            }

            //if ps meta file does not exist, just return null
            if (psMetaFilePath == null) {
                return null;
            }

            //read ps meta from file and deserialize it
            FSDataInputStream inputStream = fs.open(psMetaFilePath);
            int size = inputStream.readInt();
            Map<ParameterServerId, Integer> idToNextAttemptIndexMap = new HashMap<>(size);
            for (int i = 0; i < size; i++) {
                idToNextAttemptIndexMap.put(new ParameterServerId(inputStream.readInt()), inputStream.readInt());
            }
            inputStream.close();
            return idToNextAttemptIndexMap;
        } finally {
            psMetaLock.unlock();
        }
    }
}
